package io.kiki.sba.registry.server.metadata.slot.util.builder;

import io.kiki.sba.registry.common.model.slot.DataStoreNode;
import io.kiki.sba.registry.common.model.slot.Slot;
import io.kiki.sba.registry.common.model.slot.Slots;
import io.kiki.sba.registry.server.metadata.slot.util.MigratableSlotInfo;
import io.kiki.sba.registry.util.DatumVersionUtil;
import lombok.Getter;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.*;
import java.util.stream.Collectors;


public class SlotsBuilder implements Builder<Slots> {

    private static final Logger logger = LoggerFactory.getLogger(SlotsBuilder.class);

    private final Map<Integer, SlotBuilder> slotIdToSlotBuilderMap = new HashMap<>(256);

    private final Map<String, DataStoreNode> dataStoreNodeIdToDataStoreNodeMap = new HashMap<>();

    @Getter
    private final int slotNums;

    private final int followerNums;
    private final Slots initSlots;
    private long epoch;

    public SlotsBuilder(Slots initSlots, int slotNums, int slotReplicas) {
        this.slotNums = slotNums;
        this.followerNums = slotReplicas - 1;
        this.initSlots = initSlots;
    }

    public SlotBuilder getOrCreate(int slotId) {
        return slotIdToSlotBuilderMap.computeIfAbsent(slotId, k -> new SlotBuilder(slotId, followerNums));
    }

    public void init(List<String> dataStoreNodeIdList) {
        for (int slotId = 0; slotId < slotNums; slotId++) {
            Slot slot = initSlots == null ? null : initSlots.getSlot(slotId);
            if (slot == null) {
                getOrCreate(slotId);
                continue;
            }
            SlotBuilder slotBuilder = new SlotBuilder(slotId, slot.getLeaderDataStoreNodeId(), slot.getLeaderEpoch(), followerNums);
            if (!slotBuilder.addFollowerDataStoreNodeId(initSlots.getSlot(slotId).getFollowerDataStoreNodeIdSet())) {
                throw new IllegalArgumentException(String.format("to many followers, %s", slotBuilder));
            }
            slotIdToSlotBuilderMap.put(slotId, slotBuilder);
        }
        _reverse_(dataStoreNodeIdList);
    }

    private void _reverse_(List<String> dataStoreNodeIdList) {
        for (int slotId = 0; slotId < slotNums; slotId++) {
            SlotBuilder slotBuilder = getOrCreate(slotId);
            String leaderDataStoreNodeId = slotBuilder.getLeaderDataStoreNodeId();
            if (leaderDataStoreNodeId != null) {
                DataStoreNode dataStoreNode = dataStoreNodeIdToDataStoreNodeMap.computeIfAbsent(leaderDataStoreNodeId, k -> new DataStoreNode(leaderDataStoreNodeId));
                dataStoreNode.addLeaderSlotId(slotId);
            }
            Set<String> followerDataStoreNodeIdSet = slotBuilder.getFollowerDataStoreNodeIdSet();
            for (String followerDataStoreNodeId : followerDataStoreNodeIdSet) {
                DataStoreNode dataStoreNode = dataStoreNodeIdToDataStoreNodeMap.computeIfAbsent(followerDataStoreNodeId, k -> new DataStoreNode(followerDataStoreNodeId));
                dataStoreNode.addFollowerSlotId(slotId);
            }
        }
        /**
         * dataServer补齐
         */
        dataStoreNodeIdList.forEach(dataStoreNodeId -> dataStoreNodeIdToDataStoreNodeMap.putIfAbsent(dataStoreNodeId, new DataStoreNode(dataStoreNodeId)));
    }


    public SlotsBuilder addFollowerDataStoreNodeId(int slotId, String followerDataStoreNodeId) {
        SlotBuilder slotBuilder = getOrCreate(slotId);
        if (!slotBuilder.addFollowerDataStoreNodeId(followerDataStoreNodeId)) {
            throw new IllegalArgumentException(String.format("to many followers, %s", slotBuilder));
        }
        DataStoreNode dataStoreNode = dataStoreNodeIdToDataStoreNodeMap.computeIfAbsent(followerDataStoreNodeId, k -> new DataStoreNode(followerDataStoreNodeId));
        dataStoreNode.addFollowerSlotId(slotId);
        return this;
    }

    public SlotsBuilder removeFollowerDataStoreNodeId(int slotId, String followerDataStoreNodeId) {
        SlotBuilder slotBuilder = getOrCreate(slotId);
        if (slotBuilder.removeFollowerDataStoreNodeId(followerDataStoreNodeId)) {
            dataStoreNodeIdToDataStoreNodeMap.get(followerDataStoreNodeId).removeFollowerSlotId(slotId);
        }
        return this;
    }


    public void removeDataServerSlots(String followerDataStoreNodeId) {
        for (SlotBuilder slotBuilder : slotIdToSlotBuilderMap.values()) {
            if (slotBuilder.removeFollowerDataStoreNodeId(followerDataStoreNodeId)) {
                logger.info("[removeDataServerSlots] slot [{}] remove follower data-server[{}]", slotBuilder.getSlotId(), followerDataStoreNodeId);
            }
            if (followerDataStoreNodeId.equals(slotBuilder.getLeaderDataStoreNodeId())) {
                logger.info("[removeDataServerSlots] slot [{}] remove leader data-server[{}]", slotBuilder.getSlotId(), followerDataStoreNodeId);
                slotBuilder.setLeaderDataStoreNodeId(null);
            }
        }
        dataStoreNodeIdToDataStoreNodeMap.remove(followerDataStoreNodeId);
    }

    public String getLeaderDataStoreNodeId(int slotId) {
        SlotBuilder slotBuilder = getOrCreate(slotId);
        return slotBuilder.getLeaderDataStoreNodeId();
    }

    public List<String> getFollowerDataStoreNodeIdList(int slotId) {
        SlotBuilder slotBuilder = getOrCreate(slotId);
        return new ArrayList<>(slotBuilder.getFollowerDataStoreNodeIdSet());
    }

    public List<DataStoreNode> getDataNodeSlotsLeaderBeyond(int num) {
        return dataStoreNodeIdToDataStoreNodeMap.values().stream().filter(dataNodeSlot -> dataNodeSlot.getLeaderSlotIdSet().size() > num).collect(Collectors.toList());
    }

    public List<DataStoreNode> getDataNodeSlotsLeaderBelow(int num) {
        return dataStoreNodeIdToDataStoreNodeMap.values().stream().filter(dataNodeSlot -> dataNodeSlot.getLeaderSlotIdSet().size() < num).collect(Collectors.toList());
    }

    public List<DataStoreNode> getDataNodeSlotsFollowerBeyond(int num) {
        return dataStoreNodeIdToDataStoreNodeMap.values().stream().filter(dataNodeSlot -> dataNodeSlot.getFollowerSlotIdSet().size() > num).collect(Collectors.toList());
    }

    public List<DataStoreNode> getDataNodeSlotsFollowerBelow(int num) {
        return dataStoreNodeIdToDataStoreNodeMap.values().stream().filter(dataNodeSlot -> dataNodeSlot.getFollowerSlotIdSet().size() < num).collect(Collectors.toList());
    }

    public void incrEpoch() {
        this.epoch = DatumVersionUtil.nextId();
    }

    @Override
    public Slots build() {
        Map<Integer, Slot> map = new HashMap<>();
        slotIdToSlotBuilderMap.forEach((slotId, slotBuilder) -> {
            map.put(slotId, slotBuilder.build());
            epoch = Math.max(epoch, map.get(slotId).getLeaderEpoch());
        });
        return new Slots(epoch, map.values());
    }

    public DataStoreNode getDataStoreNode(String DataStoreNodeId) {
        DataStoreNode dataStoreNode = getDataNodeSlotIfPresent(DataStoreNodeId);
        if (dataStoreNode == null) {
            throw new IllegalArgumentException("no DataNodeSlot for " + DataStoreNodeId);
        }
        return dataStoreNode;
    }

    public DataStoreNode getDataNodeSlotIfPresent(String dataServer) {
        return dataStoreNodeIdToDataStoreNodeMap.get(dataServer);
    }


    public int getSlotReplicas() {
        return followerNums + 1;
    }

    public String replaceLeaderDataStoreNodeId(int slotId, String nextLeaderDataStoreNodeId) {
        SlotBuilder slotBuilder = getOrCreate(slotId);
        String prevLeaderDataStoreNodeId = slotBuilder.getLeaderDataStoreNodeId();
        slotBuilder.setLeaderDataStoreNodeId(nextLeaderDataStoreNodeId).removeFollowerDataStoreNodeId(nextLeaderDataStoreNodeId);
        //
        DataStoreNode nextLeaderDataStoreNode = dataStoreNodeIdToDataStoreNodeMap.computeIfAbsent(nextLeaderDataStoreNodeId, k -> new DataStoreNode(nextLeaderDataStoreNodeId));
        nextLeaderDataStoreNode.addLeaderSlotId(slotId);
        nextLeaderDataStoreNode.removeFollowerSlotId(slotId);
        //
        if (!StringUtils.isEmpty(prevLeaderDataStoreNodeId)) {
            dataStoreNodeIdToDataStoreNodeMap.get(prevLeaderDataStoreNodeId).removeLeaderSlotId(slotId);
        }
        return prevLeaderDataStoreNodeId;
    }

    /**
     * 全体的slot还可以继续分配.
     */
    public boolean hasNoAssignedSlots() {
        if (slotIdToSlotBuilderMap.size() < slotNums) {
            return true;
        }
        for (SlotBuilder slotBuilder : slotIdToSlotBuilderMap.values()) {
            if (StringUtils.isEmpty(slotBuilder.getLeaderDataStoreNodeId())) {
                return true;
            }
            if (slotBuilder.getFollowerDataStoreNodeIdSetSize() < followerNums) {
                return true;
            }
        }
        return false;
    }

    /**
     * 获取没有被分配的slot信息。
     */
    public MigratableSlotInfo getAssignableSlots() {
        MigratableSlotInfo migratableSlotInfo = new MigratableSlotInfo();
        for (int slotId = 0; slotId < slotNums; slotId++) {
            SlotBuilder slotBuilder = getOrCreate(slotId);
            if (StringUtils.isEmpty(slotBuilder.getLeaderDataStoreNodeId())) {
                migratableSlotInfo.addLeaderSlotId(slotId);
            }
            int lackFollowerNums = followerNums - slotBuilder.getFollowerDataStoreNodeIdSet().size();
            if (lackFollowerNums > 0) {
                migratableSlotInfo.addLackFollowerNum(slotId, lackFollowerNums);
            }
        }
        return migratableSlotInfo;
    }
}
